
import pymel.core as pm

def find_texture(node):
    # check if the given node is a File
    if isinstance(node, pm.nt.File):
        return node

    inputs = node.inputs()
    for input in inputs:
        if isinstance(input, pm.nt.File):
            return input
        else:
            # walk through the inputs of this node
            return find_texture(input)


def convert_material(material):
    """converts the given material to aiStandard
    """
    # get color or texture of the material
    texture = None
    color_inputs = material.diffuse.inputs()
    if color_inputs:
        texture = find_texture(color_inputs[0])

    if not texture:
        # it doesn't have a texture so it should be a gamma node
        if color_inputs:
            gamma_node = material.inputs()[0]
            if isinstance(gamma_node, pm.nodetypes.Mip_gamma_gain):
                color = gamma_node.input.get()
            elif isinstance(gamma_node, pm.nodetypes.GammaCorrect):
                color = gamma_node.value.get()
        else:
            # no gamma no connection, just get the color
            color = material.diffuse.get()

    bump_texture = None
    bump_inputs = material.standard_bump.inputs()
    if bump_inputs:
        bump_texture = find_texture(bump_inputs[0])
        bump_depth = bump_inputs[0].bumpDepth.get()

    # get other material properties
    diffuse_weight = material.diffuse_weight.get()
    diffuse_roughness = material.diffuse_roughness.get()

    reflection_color = material.refl_color.get()
    reflectivity = material.reflectivity.get()
    glossiness = material.refl_gloss.get()

    # create an arnold material
    ai = pm.shadingNode('aiStandard', asShader=1)

    # set the attributes

    # color or texture
    if texture:
        # create an aiImage node
        aiImage = pm.shadingNode('aiImage', asShader=1)
        # set the texture path
        aiImage.filename.set(texture.fileTextureName.get())
        # connect it to diffuse color
        aiImage.outColor >> ai.color
    else:
        # set the color
        ai.color.set(color)

    # bump texture
    if bump_texture:
        aiImage = pm.nodetypes.AiImage()
        aiImage.filename.set(bump_texture.fileTextureName.get())
        bumpNode = pm.nodetypes.Bump2d()
        aiImage.outAlpha >> bumpNode.bumpValue
        bumpNode.outNormal >> ai.normalCamera
        bumpNode.bumpDepth.set(bump_depth)

    # set the other attributes
    ai.Kd.set(diffuse_weight * 0.7) # normalize it by 0.7
    ai.diffuseRoughness.set(diffuse_roughness)
    ai.specularRoughness.set( 1 - glossiness)
    ai.KsColor.set(reflection_color)
    ai.Ks.set(reflectivity) # specular weight
    #ai.Kr.set(reflectivity) # reflection weight

    # assign the objects to the ai shader
    pm.hyperShade(objects=material)
    pm.hyperShade(assign=ai)

def convert_texture(texture):
    """converts the given texture to arnold texture
    """
    #
    # Convert FileTexture to aiTexture
    #
    texture_nodes = pm.ls(sl=1)
    for node in texture_nodes:
        texture_path = node.fileTextureName.get()
        conns = node.outColor.outputs(p=1)

        aiImage = pm.shadingNode('aiImage', asTexture=1)
        aiImage.filename.set(texture_path)

        for conn in conns:
            aiImage.outColor >> conn

    pm.delete(texture_nodes)
